import pandas as pd
from sklearn.metrics import mean_absolute_error
from lightgbm import LGBMRegressor
from sklearn.model_selection import GridSearchCV
import joblib  # 用于保存模型

# 读取数据
data = pd.read_csv('data.csv')

# 数据预处理（转换时间、排序、创建滞后特征、删除NaN）
data['time'] = pd.to_datetime(data['time'])
data = data.sort_values(by='time').reset_index(drop=True)
data['IC_DCS_CHIMNEY_NOX_lagged'] = data['IC_DCS_CHIMNEY_NOX'].shift(3)
data = data.dropna()

# 分离特征和标签
y = data['IC_DCS_CHIMNEY_NOX_lagged']
X = data.drop(columns=['IC_DCS_CHIMNEY_NOX', 'IC_DCS_CHIMNEY_NOX_lagged', 'time'])

# 基于时间戳划分训练集和测试集
# 我们将使用一个固定的时间点来划分，而不是随机划分
# 假设我们使用最后20%的时间段作为测试集
train_cutoff = data['time'].quantile(0.8)
X_train = X[data['time'] < train_cutoff]
y_train = y[data['time'] < train_cutoff]
X_test = X[data['time'] >= train_cutoff]
y_test = y[data['time'] >= train_cutoff]

# 初始化LightGBM回归器和参数网格
model_base = LGBMRegressor(objective='regression', metric='mae', verbose=0, random_state=42)
param_grid = {
    'num_leaves': [31, 50],
    'learning_rate': [0.1, 0.2],
    'n_estimators': [100, 200],
    'max_depth': [3, 5],
    'min_data_in_leaf': [20, 50],
    'bagging_fraction': [0.6, 0.8],
    'feature_fraction': [0.6, 0.8]
}

# 注意：LightGBM的GridSearchCV实现可能与其他库有所不同，它可能需要一个自定义的CV函数或使用`lightgbm.sklearn`中的`LGBMRegressor`配合`scikit-learn`的`GridSearchCV`。
# 这里我们直接使用`scikit-learn`的`GridSearchCV`，但请注意，对于某些参数（如`num_leaves`和`max_depth`），LightGBM有其特定的内部优化，
# 因此这些参数在网格搜索中的效果可能与在XGBoost或CatBoost中不同。

# 初始化GridSearchCV
grid_search = GridSearchCV(estimator=model_base, param_grid=param_grid, scoring='neg_mean_absolute_error', cv=3, n_jobs=-1)

# 训练模型（进行网格搜索）
grid_search.fit(X_train, y_train)

# 获取最佳参数和模型
best_params = grid_search.best_params_
best_model = grid_search.best_estimator_  # 使用GridSearchCV找到的最佳模型

# 预测测试集
y_pred = best_model.predict(X_test)

# 保存模型
joblib.dump(best_model, "xgboost_model.joblib")

# 计算MAE
mae = mean_absolute_error(y_test, y_pred)
print(f'Mean Absolute Error: {mae}')